{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20d2b3d5-0940-4744-8ba5-bc86af89ee48",
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt 在obs中的格式\n",
    "# output:\n",
    "# -    checkpoint\n",
    "# --        rank_xx\n",
    "# ---            xxx.ckpt\n",
    "# -    strategy\n",
    "# --        ckpt_strategyxxx.ckpt\n",
    "# -    log\n",
    "# 1：ckpt瘦身\n",
    "# 2：如何有pp模式 需要将strategy合并\n",
    "# 3：将ckpt合并"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6d34455-1971-4899-9f42-bede44b05d98",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a408e22f-078e-4342-989a-e7e3cc79a8fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#传送至notebook\n",
    "import moxing as mox\n",
    "mox.file.copy_parallel('obs://hw-test/output/strategy/ziya_4_141/strategy/', '/home/ma-user/work/ziya/4/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af2e2f43-5720-44cb-a53f-b5c967f4e8a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#合并strategy权重\n",
    "import mindspore as ms\n",
    "ms.merge_pipeline_strategys(\"源文件夹所在目录\",\"需要生成的文件名称+路径\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acf01f48-6bd2-41ef-ac8a-f3e36273ed35",
   "metadata": {},
   "outputs": [],
   "source": [
    "#权重瘦身\n",
    "import argparse\n",
    "from mindspore import load_checkpoint, save_checkpoint\n",
    "import copy\n",
    "\n",
    "\n",
    "def print_dict(input_dict):\n",
    "    \"\"\"\n",
    "    Print the keys and values of input dict\n",
    "\n",
    "    Args:\n",
    "        input_dict(dict): input dict with key and value.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    for k, v in input_dict.items():\n",
    "        print(f\"Param: {k} with shape {v}\")\n",
    "\n",
    "\n",
    "def get_converted_ckpt(weight_dict):\n",
    "    new_ckpt_list = []\n",
    "    for src, value in weight_dict.items():\n",
    "        new_ckpt_list.append({\"data\": value, \"name\": src})\n",
    "    return new_ckpt_list\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    #32卡为例 多开几个for循环加速 \n",
    "    for i in range(32):\n",
    "        #mindspore_path = '/home/ma-user/work/ziya/ckpt_pretrain/rank_'+str(i)+'/llama_13b_rank_'+str(i)+'-3298_2.ckpt'\n",
    "        #slimmed_mindspore_path = '/home/ma-user/work/ziya/new_ckpt_pretrain/rank_'+str(i)+'/new_llama_13b_rank_'+str(i)+'-3298_2.ckpt'\n",
    "        mindspore_path = '需要瘦身的src_ckpt'\n",
    "        slimmed_mindspore_path = '瘦身后的dst_ckpt'\n",
    "        state_dict = load_checkpoint(mindspore_path)\n",
    "        print_dict(state_dict)\n",
    "        #删除不需要的权重\n",
    "        not_need_parameters = [\"adam\", \"scale_sense\", \"global_step\", \"current_iterator_step\", \"last_overflow_iterator_step\",\n",
    "                               \"epoch_num\", \"step_num\", \"loss_scale\",\"accu_grads\"]\n",
    "\n",
    "        keys_list = copy.deepcopy(list(state_dict.keys()))\n",
    "        for _key in keys_list:\n",
    "            for _name in not_need_parameters:\n",
    "                if _key.startswith(_name):\n",
    "                    del state_dict[_key]\n",
    "                    break\n",
    "\n",
    "        save_checkpoint(get_converted_ckpt(state_dict), slimmed_mindspore_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b67955da-267d-473d-979b-08684e38405c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#权重合并\n",
    "from mindspore import transform_checkpoints\n",
    "transform_checkpoints(\n",
    "    src_checkpoints_dir=\"/home/ma-user/work/ziya/new_ckpt_pretrain/\", # 原切分权重文件夹\n",
    "    dst_checkpoints_dir=\"/home/ma-user/work/ziya/1_ckpt_pretrain_test/\", # 目标路径\n",
    "    ckpt_prefix=\"ziya-13b-pretrain\", # .ckpt文件前缀名\n",
    "    src_strategy_file=\"/home/ma-user/work/ziya/output2/new_strategy.ckpt\", # 步骤1中的切分策略文件路径\n",
    "    dst_strategy_file=None # None表示不切分，权重合一\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eef54e5-246c-474a-bc29-c12a19c32cb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MindSpore",
   "language": "python",
   "name": "mindspore"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
